import os
"""
 * @Author       : Benjay·Shaw
 * @Date         : 2021-08-11 20:13:03
 * @LastEditors  : Benjay·Shaw
 * @LastEditTime : 2021-08-27 12:40:36
 * @Description  : 存储点位经纬度的shapefile
"""
import shapefile
from osgeo import osr, ogr
from tqdm import tqdm
"""
 * @description: 设置shp文件坐标系
 * @param {*} shp_path
 * @param {*} out_shp_path
 * @return {*}
"""


def add_SpatialReference(shp_path, out_shp_path):
    driver = ogr.GetDriverByName('ESRI Shapefile')
    dataSource = driver.Open(shp_path, 1)
    layer = dataSource.GetLayer()
    LayerDefn = layer.GetLayerDefn()
    SpatialReference = osr.SpatialReference()
    SpatialReference.ImportFromEPSG(4326)
    out_DataSource = driver.CreateDataSource(out_shp_path)
    out_Layer = out_DataSource.CreateLayer(out_shp_path, SpatialReference,
        ogr.wkbPoint)
    def_feature = out_Layer.GetLayerDefn()
    for i in range(0, LayerDefn.GetFieldCount()):
        fieldDefn = LayerDefn.GetFieldDefn(i)
        out_Layer.CreateField(fieldDefn)
    for feature in layer:
        geom = feature.GetGeometryRef()
        out_feature = ogr.Feature(def_feature)
        for i in range(0, LayerDefn.GetFieldCount()):
            out_feature.SetField(LayerDefn.GetFieldDefn(i).GetNameRef(),
                feature.GetField(i))
        out_feature.SetGeometry(geom)
        out_Layer.CreateFeature(out_feature)
    out_DataSource.FlushCache()
    del dataSource
    SpatialReference.MorphFromESRI()
    prjfile = open(shp_path.replace('.shp', '.prj'), 'w')
    prjfile.write(SpatialReference.ExportToWkt())
    prjfile.close()


cur_path = os.path.dirname(__file__)
shp = shapefile.Writer(target=cur_path + '/input/PointInfo.shp', shapeType=
    shapefile.POINT, autoBalance=True)
shp.field('Longitude', 'C', '128', 0)
shp.field('Latitude', 'C', '128', 0)
shp.field('Confidence', 'C', '128', 0)
files = os.listdir(cur_path + '/input/detection-latlon')
loop = tqdm(enumerate(files), total=len(files))
for _, file in loop:
    position = cur_path + '/input/detection-latlon/' + file
    loop.set_postfix(file=file)
    if os.path.getsize(position) > 0:
        with open(position, 'r', encoding='utf-8') as f:
            for line in f:
                if len(line.strip().split()) == 3:
                    longitude = line.strip().split()[0]
                    latitude = line.strip().split(' ')[1]
                    confidence = line.strip().split(' ')[2]
                    shp.point(float(longitude), float(latitude))
                    shp.record(longitude, latitude, confidence)
shp.close()
shp_path = cur_path + '/input/PointInfo.shp'
out_shp_path = cur_path + '/input/PointInfo_0.8_0.7.shp'
add_SpatialReference(shp_path, out_shp_path)
